package cn.lijiancn01.demo.system.codegen;

import cn.lijiancn01.demo.framework.exception.BizException;
import com.mybatisflex.codegen.config.GlobalConfig;
import com.mybatisflex.codegen.config.StrategyConfig;
import com.mybatisflex.codegen.dialect.IDialect;
import com.mybatisflex.codegen.entity.Table;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.Marker;
import org.slf4j.MarkerFactory;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.DatabaseMetaData;
import java.sql.ResultSet;
import java.util.ArrayList;
import java.util.List;

@Slf4j
public class TableUtils {
    public static final Marker marker = MarkerFactory.getMarker("TableUtils");

    public static List<Table> buildTables(DataSource dataSource, GlobalConfig globalConfig) {
        return buildTables(dataSource, globalConfig, null);
    }

    public static List<Table> buildTables(DataSource dataSource, GlobalConfig globalConfig, IDialect dialect) {
        try (Connection conn = dataSource.getConnection()) {
            if (dialect == null) {
                dialect = IDialect.DEFAULT;
            }
            DatabaseMetaData dbMeta = conn.getMetaData();
            StrategyConfig strategyConfig = globalConfig.getStrategyConfig();
            String schemaName = strategyConfig.getGenerateSchema();
            List<Table> tables = new ArrayList<>();
            try (ResultSet rs = getTablesResultSet(conn, globalConfig, dialect)) {
                while (rs.next()) {
                    String tableName = rs.getString("TABLE_NAME");
                    //黑名单
                    if (!strategyConfig.isSupportGenerate(tableName)) {
                        continue;
                    }

                    Table table = new Table();
                    table.setGlobalConfig(globalConfig);
                    table.setTableConfig(strategyConfig.getTableConfig(tableName));

                    table.setSchema(schemaName);
                    table.setName(tableName);

                    String remarks = rs.getString("REMARKS");
                    table.setComment(remarks);

                    buildPrimaryKey(conn, table);

                    dialect.buildTableColumns(schemaName, table, globalConfig, dbMeta, conn);

                    tables.add(table);
                }
            }
            return tables;
        } catch (Exception e) {
            log.error(marker, "buildTables error", e);
            throw new BizException(e.getMessage());
        }
    }

    public static ResultSet getTablesResultSet(Connection conn, GlobalConfig globalConfig, IDialect dialect) {
        try {
            String schema = globalConfig.getStrategyConfig().getGenerateSchema();
            DatabaseMetaData dbMeta = conn.getMetaData();
            if (globalConfig.getStrategyConfig().isGenerateForView()) {
                return dialect.getTablesResultSet(dbMeta, conn, schema, new String[]{"TABLE", "VIEW"});
            } else {
                return dialect.getTablesResultSet(dbMeta, conn, schema, new String[]{"TABLE"});
            }
        } catch (Exception e) {
            log.error(marker, "getTablesResultSet error", e);
            throw new BizException(e.getMessage());
        }
    }

    public static void buildPrimaryKey(Connection conn, Table table) {
        try {
            DatabaseMetaData dbMeta = conn.getMetaData();
            try (ResultSet rs = dbMeta.getPrimaryKeys(conn.getCatalog(), null, table.getName())) {
                while (rs.next()) {
                    String primaryKey = rs.getString("COLUMN_NAME");
                    table.addPrimaryKey(primaryKey);
                }
            }
        } catch (Exception e) {
            log.error(marker, "buildPrimaryKey error", e);
            throw new BizException(e.getMessage());
        }
    }
}
